from typing import TypedDict
import requests
import logging

class NovelState(TypedDict):
    current_outline: str
    chapters: dict
    iteration_count: int

class NovelNodes:
    def __init__(self):
        self.llm_endpoint = 'https://open.bigmodel.cn/api/paas/v4/chat/completions'
        self.llm_key = '734c6ee4a5524052a1fc44580660dc38.F5A02rjHWxocAsaT'
        
    def _call_llm(self, prompt: str) -> str:
        logging.info(f"调用LLM，提示词：{prompt[:50]}...")
        headers = {
            'Authorization': f'Bearer {self.llm_key}',
            'Content-Type': 'application/json'
        }
        data = {
            "model": "glm-4-flash",
            "messages": [{"role": "user", "content": prompt}]
        }
        response = requests.post(self.llm_endpoint, json=data, headers=headers)
        return response.json()['choices'][0]['message']['content']

    def generate_initial_outline(self, state: NovelState) -> NovelState:
        prompt = "请生成一个小说的大纲，包含主要章节和情节发展。"
        state['current_outline'] = self._call_llm(prompt)
        return state

    def refine_outline(self, state: NovelState) -> NovelState:
        prompt = f"当前大纲：{state['current_outline']}\n请改进此小说大纲。"
        state['current_outline'] = self._call_llm(prompt)
        state['iteration_count'] += 1
        return state

    def write_chapter(self, state: NovelState, chapter_name: str) -> NovelState:
        prompt = f"大纲参考：{state['current_outline']}\n请撰写章节：{chapter_name}，控制在500字左右。"
        content = self._call_llm(prompt)[:50]
        state['chapters'][chapter_name] = content
        logging.info(f"章节 {chapter_name} 生成完成，字数：{len(content)}")
        return state

    def final_review(self, state: NovelState) -> NovelState:
        prompt = f"请对完整小说进行最终审核：{state['chapters']}"
        state['review'] = self._call_llm(prompt)
        return state